/*
 * @(#) $Header$
 *
 * Copyright (C) 2007  Forklabs Daniel Léonard
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
 */

package ca.forklabs.media.jai.opimage;

import java.awt.Point;
import java.awt.Rectangle;
import java.awt.image.Raster;
import java.awt.image.RenderedImage;
import java.awt.image.WritableRaster;
import java.util.Arrays;
import java.util.Comparator;
import java.util.Map;
import java.util.Random;
import javax.media.jai.ImageLayout;
import javax.media.jai.OpImage;
import javax.media.jai.UntiledOpImage;
import javax.media.jai.iterator.RandomIter;
import javax.media.jai.iterator.RandomIterFactory;
import javax.media.jai.iterator.WritableRandomIter;
import ca.forklabs.media.jai.operator.KMeansDescriptor;

/**
 * Class {@code KMeansOpImage} is the {@link OpImage} for operator
 * <em>kmeans</em>.
 *
 * @author   <a href="mailto:forklabs at gmail.com?subject=ca.forklabs.media.jai.opimage.KMeansOpImage">Daniel Léonard</a>
 * @version $Revision$
 */
@SuppressWarnings("unchecked")
public class KMeansOpImage extends UntiledOpImage {

//---------------------------
// Instance variables
//---------------------------

   /** The number of clusters to make. */
   protected int clusters;

   /** The evaluation function. */
   protected KMeansDescriptor.EvaluationFunction function;

   /** The number of iteration to perform. */
   protected int iterations;

   /** The color map of the segmented image. */
   protected int[][] color_map;


//---------------------------
// Constructors
//---------------------------

   /**
    * Constructor.
    * @param   source   the image to segment.
    * @param   clusters   the number of desired clusters.
    * @param   function   the evaluation function.
    * @param   iterations   the maximum number of iterations.
    * @param   color_map   the color_map of the segmented image.
    * @param   hints   rendering hints.
    * @param   layout  image layout.
    */
   public KMeansOpImage(RenderedImage source, int clusters, KMeansDescriptor.EvaluationFunction function, int iterations, int[][] color_map, Map<?, ?> hints, ImageLayout layout) {
      super(source, hints, layout);
      this.setup(clusters, function, iterations, color_map);
      }


//---------------------------
// Accessors and mutators
//---------------------------

   /**
    * Changes the number of clusters.
    * @param   clusters   the new number of clusters.
    */
   protected void setClusters(int clusters) {
      this.clusters = clusters;
      }

   /**
    * Gets the desired number of clusters.
    * @return   the number of clusters.
    */
   protected int getClusters() {
      return this.clusters;
      }

   /**
    * Changes the evaluation function.
    * @param   function   the new function.
    */
   protected void setEvaluationFunction(KMeansDescriptor.EvaluationFunction function) {
      this.function = function;
      }

   /**
    * Gets the evaluation function.
    * @return   the function.
    */
   protected KMeansDescriptor.EvaluationFunction getEvaluationFunction() {
      return this.function;
      }

   /**
    * Changes the number of iterations.
    * @param   iterations   the new number of iterations.
    */
   protected void setIterations(int iterations) {
      this.iterations = iterations;
      }

   /**
    * Gets the desired number of iterations.
    * @return   the number of iterations.
    */
   protected int getIterations() {
      return this.iterations;
      }

   /**
    * Changes the color map.
    * @param   color_map   the new color map.
    */
   protected void setColorMap(int[][] color_map) {
      this.color_map = color_map;
      }

   /**
    * Gets the color map.
    * @return   the color map.
    */
   protected int[][] getColorMap() {
      return this.color_map;
      }


//---------------------------
// Instance methods
//---------------------------

   /**
    * Performs a stable in-place sort of the center from smallest to greatest. A
    * center is smaller if its square distance from the origin is smaller then
    * another square distance.
    * @param   centers   the center to sort.
    */
   protected void sortCenters(double[][] centers) {
      Arrays.sort(centers, new Comparator<double[]>() {
         private double squareDistance(double[] array) {
            double distance = 0.0;
            for (int i = 0, len = array.length; i < len; i++) {
               distance += array[i] * array[i];
               }
            return distance;
            }

         @Override
         public int compare(double[] array1, double[] array2) {
            double square_len_1 = this.squareDistance(array1);
            double square_len_2 = this.squareDistance(array2);
            int comparison = (int) (square_len_1 - square_len_2);
            return comparison;
            }
         });
      }

   /**
    * Initialize the new centers. Each dimension of each center is selected
    * randomly between <em>[0.0, 1.0[</em> and strive to be different from
    * each other.
    * @param   clusters   the number of clusters.
    * @param   dimensions   the dimension of each cluster.
    * @return   the initial centers.
    */
   @SuppressWarnings("hiding")
   protected double[][] initializeCenters(int clusters, int dimensions) {
      double[][] centers = new double[clusters][dimensions];

   // the centers are selected at random, but they all must be different
      Random random = new Random();
      for (int c = 0, tries = 0; (c < clusters); ) {

         double[] current = centers[c];

         for (int d = 0; d < dimensions; d++) {
            current[d] = random.nextDouble();
            }

         boolean center_already_exist = false;
         for (int i = 0; (false == center_already_exist) && (i < c); i++) {
            double[] previous = centers[i];
            center_already_exist = Arrays.equals(current, previous);
            }

      // if the center already exist, we try again,
      // but if we tried too much, we take what we
      // have and continue
         if ((center_already_exist && (tries < 100))) {
            tries++;
            }
         else {
            c++;
            tries = 0;
            }
         }

      this.sortCenters(centers);

      return centers;
      }

   /**
    * Resets the <em>new center</em> data structures.
    * @param   positions   the position of the new centers.
    * @param   population   the population of each new center.
    */
   protected void resetNewCenters(double[][] positions, double[] population) {
      for (double[] position : positions) {
         Arrays.fill(position, 0.0);
         }
      Arrays.fill(population, 0.0);
      }

   /**
    * Sets up the min vector for the normalization.
    * @param   len   the length of the vector.
    * @return   the min vector.
    */
   protected double[] setupMinVector(int len) {
      double[] min = new double[len];
      Arrays.fill(min, Double.MAX_VALUE);
      return min;
      }

   /**
    * Sets up the max vector for the normalization.
    * @param   len   the length of the vector.
    * @return   the max vector.
    */
   protected double[] setupMaxVector(int len) {
      double[] max = new double[len];
      Arrays.fill(max, -Double.MAX_VALUE);
      return max;
      }

   /**
    * Updates the min and max vector from the data in the specified vector.
    * @param   vector   the vector to look into for a new min or a new max.
    * @param   min   the min vector.
    * @param   max   the max vector.
    */
   protected void updateMinMax(double[] vector, double[] min, double[] max) {
      for (int i = 0, len = vector.length; i < len; i++) {
         min[i] = Math.min(vector[i], min[i]);
         max[i] = Math.max(vector[i], max[i]);
         }
      }

   /**
    * Normalizes the vector.
    * @param   vectors   the vectors to normalize.
    */
   protected void normalize(double[][] vectors) {
      int dim = vectors[0].length;

      double[] mins = this.setupMinVector(dim);
      double[] maxs = this.setupMaxVector(dim);

      for (int i = 0, len = vectors.length; i < len; i++) {
         this.updateMinMax(vectors[i], mins, maxs);
         }

      for (int d = 0; d < dim; d++) {
         double offset = mins[d];
         double scale = maxs[d] - mins[d];
         for (int i = 0; i < vectors.length; i++) {
            vectors[i][d] = (vectors[i][d] - offset) / scale;
            }
         }
      }

   /**
    * Calculate the square of the distance of the pixel from each center pixel.
    * @param   pixel   the pixel.
    * @param   centers   the center pixels.
    * @param   distances   the arrays to store the square distances.
    * @return   {@code distances}.
    */
   protected double[] calculateSquareDistances(double[] pixel, double[][] centers, double[] distances) {
      int num_clusters = centers.length;
      int num_bands = pixel.length;

      for (int c = 0; c < num_clusters; c++) {
         double distance = 0.0;
         for (int b = 0; b < num_bands; b++) {
            double diff = pixel[b] - centers[c][b];
            distance += diff * diff;
            }
         distances[c] = distance;
         }
      return distances;
      }

   /**
    * Gets the position of the smallest distance.
    * @param   distances   the distances.
    * @return   the position.
    */
   protected int findClosestCenter(double[] distances) {
      int cluster = 0;
      double min = distances[cluster];

      for (int c = 1, len = distances.length; c < len; c++) {
         double candidate = distances[c];
         boolean is_smaller_than_minimum = (candidate < min);
         if (is_smaller_than_minimum) {
            cluster = c;
            min = candidate;
            }
         }

      return cluster;
      }

   /**
    * Gets the error message telling that the raster data type is unknown.
    * @param   type   the bad type.
    * @return   the formatted error message.
    */
   @SuppressWarnings("boxing")
   protected String getUnknownDataTypeErrorMessage(int type) {
      String key = Resources.UNKNOWN_DATA_TYPE;
      String message = Resources.getLocalizedString(key, type);
      return message;
      }

   /**
    * Sets up this image.
    * @param   clusters   the number of desired clusters.
    * @param   function   the evaluation function.
    * @param   iterations   the number of desired iterations.
    * @param   color_map   the color map.
    */
   protected void setup(int clusters, KMeansDescriptor.EvaluationFunction function, int iterations, int[][] color_map) {
      this.setClusters(clusters);
      this.setEvaluationFunction(function);
      this.setIterations(iterations);
      this.setColorMap(color_map);
      }


//---------------------------
// Overridden methods from javax.media.jai.UntiledOpImage
//---------------------------

   /**
    * Computes the k-means image.
    * @param   sources    the source rasters.
    * @param   sink   the sink raster.
    * @param   bounds   the bounds to work within.
    */
   @Override
   @SuppressWarnings("hiding")
   protected void computeImage(Raster[] sources, WritableRaster sink, Rectangle bounds) {
      Raster source = sources[0];

   // for each iteration until stability :
   // 1) sort and reset the centers - the centers are already normalized
   // 2) for each row
   //      for each column
   //         store the pixel vector
   // 3) normalize each pixel vector
   // 4) for each pixel vector
   //   4a) calculate its distance from each center
   //   4b) find the center that is closest
   //   4c) update the data structure for the next iteration centers
   //   4d) set the pixel value of the sink to the cluster number and
   //       record if any change occured
   // 5) calculate the position of the new centers

      int width = (int) bounds.getWidth();
      int height = (int) bounds.getHeight();

      int min_x = (int) bounds.getX();
      int max_x = min_x + width;
      int min_y = (int) bounds.getY();
      int max_y = min_y + height;

      int bands = source.getNumBands();

      KMeansDescriptor.EvaluationFunction function = this.getEvaluationFunction();
      Point position = new Point();

      RandomIter source_iter = RandomIterFactory.create(source, source.getBounds());
      WritableRandomIter sink_iter = RandomIterFactory.createWritable(sink, sink.getBounds());

      int clusters = this.getClusters();
      int dim;
      double[][] vectors;
      double[][] centers;
      double[][] new_centers_position;
      {
         position.x = min_x;
         position.y = min_y;
         double[] vector = function.invoke(source_iter, position);
         dim = vector.length;

         int len = width * height;
         vectors = new double[len][dim];

         centers = this.initializeCenters(clusters, dim);
         new_centers_position = new double[clusters][dim];
      }
      double[] new_centers_population = new double[clusters];

      int iterations = this.getIterations();
      boolean is_stable = false;
      for (int i = 0; (false == is_stable) && (i < iterations); i++) {
      // 1)
         this.sortCenters(centers);
         this.resetNewCenters(new_centers_position, new_centers_population);
         is_stable = true;

      // 2)
         for (int y = min_y, v = 0; y < max_y; y++) {
            for (int x = min_x; x < max_x; x++, v++) {
               position.x = x;
               position.y = y;
               double[] vector = function.invoke(source_iter, position);
               System.arraycopy(vector, 0, vectors[v], 0, dim);
               }
            }

      // 3)
         this.normalize(vectors);

      // 4)
         double[] sink_pixel = new double[bands];
         for (int y = min_y, v = 0; y < max_y; y++) {
            for (int x = min_x; x < max_x; x++, v++) {
               double[] vector = vectors[v];
            // 4a)
               double[] distances = new double[clusters];
               distances = this.calculateSquareDistances(vector, centers, distances);
            // 4b)
               int cluster = this.findClosestCenter(distances);
            // 4c)
               for (int d = 0; d < dim; d++) {
                   new_centers_position[cluster][d] += vector[d];
                   }
               new_centers_population[cluster]++;
            // 4d)
               sink_iter.getPixel(x, y, sink_pixel);
            // For a single pixel in the sink image, all its component have the
            // same value. To determine if the pixel needs to have its color
            // changed, we only need to check the first element.
               boolean has_this_pixel_changed = (cluster != sink_pixel[0]);
               if (has_this_pixel_changed) {
                  Arrays.fill(sink_pixel, cluster);
                  sink_iter.setPixel(x, y, sink_pixel);
                  }
               is_stable &= !has_this_pixel_changed;
               }
            }

      // 5)
         for (int c = 0; c < clusters; c++) {
            for (int d = 0; d < dim; d++) {
// BUG : what happens if a center has no population ?
               centers[c][d] = new_centers_position[c][d] / new_centers_population[c];
               }
            }
         }

      double[] sink_pixel = new double[bands];
      int[][] color_map = this.getColorMap();
      for (int y = min_y; y < max_y; y++) {
         for (int x = min_x; x < max_x; x++) {
            sink_iter.getPixel(x, y, sink_pixel);

            int cluster = (int) sink_pixel[0];
            int[] color = color_map[cluster];

            sink_iter.setPixel(x, y, color);
            }
         }

      source_iter.done();
      sink_iter.done();

// BUG : set the color model as an indexed color model
      }

   }
